# 代码编写者： 李朔
# 开发时间： 2021/7/11 11:56
import torch
import matplotlib.pyplot as plt
#训练数据集
x_data=[1.0,2.0,3.0]
y_data=[2.0,4.0,6.0]

#利用Tensor变量来存放我们的权重以及偏置
w1=torch.Tensor([1.0])
w2=torch.Tensor([1.0])
b=torch.Tensor([0.0])
#计算梯度
w1.requires_grad=True
w2.requires_grad=True
b.requires_grad=True

#model
def forward(x):
    return w1*x*x+w2*x+b

#损失函数
def loss(x,y):
    y_pred=forward(x)
    return (y_pred-y)**2


#训练过程
loss_list=[]
print("predict (before training)",4,forward(4).item())
for epoch in range(100):
    for x,y in zip(x_data,y_data):
        l=loss(x,y)
        l.backward()
        print("\tgrad:(x={0},y={1})".format(x,y)+"w1.grad={0},w2.grad={1},b.grad={2}".format(w1.grad.item(),w2.grad.item(),b.grad.item()))
        w1.data=w1.data-0.01*w1.grad.data
        w2.data = w2.data - 0.01 * w2.grad.data
        b.data = b.data - 0.01 * b.grad.data
        w1.grad.data.zero_()
        w2.grad.data.zero_()
        b.grad.data.zero_()
    loss_list.append(l.item())
    print("porgress:{0}\t loss={1}".format(epoch,l.item()))
print("predict (after training)",4,forward(4).item())

#绘图
plt.plot(loss_list)
plt.xlabel("epoch")
plt.ylabel("loss")
plt.show()